'''
模型训练类
'''

import joblib as jb
from pandas import DataFrame
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.multioutput import MultiOutputRegressor, RegressorChain
from sklearn.tree import DecisionTreeRegressor
from xgboost import XGBRegressor

from config.global_config import Config
from log.LogModel import LogModel
from predict.predict_ai import test_model_predict, xgb_importance_plot, Multi_importance_plot
from prepare.feature_engineer import feature_engineer, feature_engineer_origin
from sklearn.preprocessing import StandardScaler


class ModelTrain(object):
    '''
    随机森林 网格搜索交叉验证
    '''
    def random_forest_grid_search(self, load_config_model, x_train, y_train):
        # 随机森林
        rf = RandomForestRegressor(load_config_model.grid_search_verbose, max_depth=load_config_model.max_depth)
        rf.fit(x_train, y_train)
        # 网格搜索交叉验证
        params = {'n_estimators': load_config_model.grid_search_n_estimators,
                  'max_depth': load_config_model.grid_search_max_depth,
                  'random_state': load_config_model.grid_search_random_state}
        model = GridSearchCV(estimator=rf, param_grid=params, cv=load_config_model.grid_search_cv)
        model.fit(x_train, y_train)
        print(f'随机森林网格搜索后的最佳参数: {model.best_estimator_}')

    '''
    XGB 回归模型训练
    '''

    def xgb_regressor_model_train(self,
                                  x: DataFrame,
                                  y: DataFrame,
                                  logger):
        logger.info('--------开始xgboost模型训练---------')
        load_config_model = Config('../config/xgb_regressor_model_train.txt')
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=load_config_model.test_size,
                                                            random_state=load_config_model.random_state)

        params = {
            'n_estimators': [10,20,],
            'max_depth': [2,3],
            'random_state': [8],
            'learning_rate': [0.01,0.02]
        }
        # 模型实例化
        estimator = XGBRegressor()
        estimator = GridSearchCV(estimator=estimator, param_grid=params,
                                 cv=5,scoring='neg_mean_squared_error')
        # 模型训练
        estimator.fit(x_train, y_train)

        logger.info(estimator.best_params_)
        # 模型保存
        file_name = '../model/xgb_test.pth'
        jb.dump(estimator.best_estimator_, file_name)

        return x_test, y_test, file_name



    '''
    随机森林 模型训练
    :param data 数据列表
    :param feature_names 特征
    :param label 标签
    '''

    def random_forest_model_train(self,
                                  x: DataFrame,
                                  y: DataFrame
                                  ):
        load_config_model = Config('../config/random_forest_model_train.txt')
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=load_config_model.test_size,
                                                            random_state=load_config_model.random_state)

        # 模型实例化
        estimator = RandomForestRegressor(n_estimators=load_config_model.n_estimators,
                                          max_depth=load_config_model.max_depth,
                                          random_state=load_config_model.random_state)
        # 模型训练
        estimator.fit(x_train, y_train)

        # 模型保存
        file_name = '../model/random_test.pth'
        jb.dump(estimator, file_name)

        # 随机森林 网格搜索交叉验证
        self.random_forest_grid_search(load_config_model, x_train, y_train)
        return x_test, y_test, file_name

    '''
    多输出回归器 模型训练
    :param data 数据列表
    :param feature_names 特征
    :param label 标签
    '''

    def multi_output_regressor_model_train(self,
                                           x: DataFrame,
                                           y: DataFrame,
                                           logger):
        logger.info('--------开始多输出模型训练---------')
        load_config_model = Config('../config/multi_output_regressor_model_train.txt')
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=load_config_model.test_size,
                                                            random_state=load_config_model.random_state)
        scaler = StandardScaler()
        x_train_s = scaler.fit_transform(x_train)
        x_test_s = scaler.transform(x_test)

        params = {
            'estimator__n_estimators': [10,20,30,40,50],
            'estimator__max_depth': [2],
            'estimator__random_state': [8],
            'estimator__learning_rate': [0.01]
        }
        # 模型实例化
        estimator = XGBRegressor()
        estimator = MultiOutputRegressor(estimator=estimator)
        estimator = GridSearchCV(estimator=estimator, param_grid=params,
                                 cv=5, scoring='r2')
        # 模型训练
        estimator.fit(x_train_s, y_train)

        logger.info(estimator.best_params_)
        # 模型保存
        file_name = '../model/multi_test.pth'
        jb.dump(estimator.best_estimator_, file_name)

        return x_test_s, y_test, file_name, x_test.columns

    '''
     链式回归 模型训练
    :param data 数据列表
    :param feature_names 特征
    :param label 标签
    '''

    def chained_regressor_model_train(self,
                                      x: DataFrame,
                                      y: DataFrame
                                      ):
        load_config_model = Config('../config/chained_regressor_model_train.txt')
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=load_config_model.test_size,
                                                            random_state=load_config_model.random_state)

        # 模型实例化
        estimator = RegressorChain(base_estimator=LinearRegression(), random_state=load_config_model.random_state)
        # 模型训练
        estimator.fit(x_train, y_train)

        # 模型保存
        file_name = '../model/chained_test.pth'
        jb.dump(estimator, file_name)

        return x_train, y_train, x_test, y_test, file_name


if __name__ == '__main__':
    data_path = '../data/新数据成长性_合并.xlsx'
    log_model = LogModel()
    model = ModelTrain()

    x, y = feature_engineer(log_model.data_preprocess, log_model.logfile)

    x_test, y_test, model_path = model.xgb_regressor_model_train(x, y, log_model.logfile)

    test_model_predict(x_test, y_test, model_path, log_model.logfile)

    xgb_importance_plot(model_path, log_model.logfile)

    x_test, y_test, model_path, feature_names = model.multi_output_regressor_model_train(x, y, log_model.logfile)

    test_model_predict(x_test, y_test, model_path, log_model.logfile)

    Multi_importance_plot(model_path, feature_names, y_test.columns, log_model.logfile)

